import pandas as pd
import numpy as np
import os


# 读取文件
def read_code(code, year):
    data = pd.read_csv('./stocks/%s/%s.csv' % (year, code))
    data[code] = code
    return data


def read_all(code):
    return pd.concat([read_code(code, str(i)) for i in range(2013, 2021)])


def count_high(df):
    # 先扩展列
    if df.empty:
        return
    first_date, = df.head(1)['date']

    new_data = df.assign(temp=9.5).assign(sign=lambda x: x['temp'] < x['pctChg']).drop('temp', axis=1)
    pos, = np.where(np.diff(new_data['sign']))
    start = np.insert(pos + 1, 0, 0)
    end = np.append(pos, len(new_data) - 1)
    tmp = pd.DataFrame({'code': list(new_data.code.iloc[start]),
                        'start_date': list(new_data.date.iloc[start]),
                        'end_date': list(new_data.date.iloc[end]),
                        'sign': list(new_data.sign.iloc[start]),
                        'start_open': list(new_data.open.iloc[start]),
                        'end_close': list(new_data.close.iloc[end]),
                        'continue_day': end - start + 1
                        })
    # 后续五天的数据通过窗口函数去拿？
    # 以及新股的判断
    tmp = tmp.assign(first_date=first_date) \
        .loc[lambda x: x['sign']] \
        .drop('sign', axis=1) \
        .loc[lambda x: x['continue_day'] > 6]
    return tmp


def check_after_high(df):
    high = count_high(df)
    if high is None:
        return None
    end_dates = high['end_date'].values
    after_raise = []
    after_raise_header = ['day_%d' % i for i in range(1, 13)]
    for end_date in end_dates:
        part = df.loc[lambda x: x['date'] > end_date]
        end_close = high[lambda x: x['end_date'] == end_date]['end_close'].values[0]
        p = part.assign(raise_up=lambda x: round(100 * (x['close'] / end_close - 1), 2))
        _np = p.head(12).raise_up.tolist()
        _np.extend([pd.NaT for i in range(0, 12 - len(_np))])
        after_raise.append(_np)
    if len(after_raise) == 0:
        return None
    aft = pd.DataFrame(after_raise, columns=after_raise_header, index=high.index)
    aft['end_date'] = end_dates
    print(high)
    print(aft)
    return pd.concat([high, aft], axis=1,  join='inner')


def _map():
    file_path = 'stocks/2018'
    datas = []
    for top, d, files in os.walk(file_path):
        for file in files:
            stock = file[:-4]
            data = check_after_high(read_all(stock))
            print('complete %s' % stock)
            if data is None or data.empty:
                continue
            datas.append(data)
    tot = pd.concat(datas)
    print(tot)
    tot.to_excel('consume/consume2013-continue.xlsx', index=False)


def reduce():
    # 合并所有map

    file_path = 'alys'
    for top, d, files in os.walk(file_path):
        if len(d) != 0:
            continue
        i = 0
        total = len(files)
        datas = []
        year = top[-4:]

        for file in files:
            print('开始 %s %s' % (top, file))
            stock = file[:-4]
            year = top[-4:]
            i += 1
            data = pd.read_csv('%s/%s' % (top, file), index_col='start_date')[lambda x: x.continue_day >= 5]
            if data.empty or len(data) == 0:
                continue
            data['code'] = stock
            datas.append(data)
            print('完成了 %s, %d/%d' % (top, i, total))
        pd.concat(datas).to_csv('consume/aly_%s.csv' % year)


if __name__ == '__main__':
    # print(check_after_high(read_all('sz.002565')))
    _map()
    # print(pd.DataFrame([[-2.35, -8.59, -11.22, -8.73, -7.2, -13.02, -17.17, -18.84, -19.39, -21.47, -22.02, -22.16]
    #                     ,[-2.35, -8.59, -11.22, -8.73, -7.2, -13.02, -17.17, -18.84, -19.39, -21.47, -22.02, -22.16],
    #                     [-2.35, -8.59, -11.22, -8.73, -7.2, -13.02, -17.17, -18.84, -19.39, -21.47, -22.02, -22.16]],
    #                    columns=['day%d' % i for i in range(1, 13)]))
    # print(pd.DataFrame([[1,3,4],[1,2]]))
    # [1,2].extend([None])
